# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2024-07-30
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-07-30

# @Description: graph database service

from utils import status

from config.settings import *

from neo4j import GraphDatabase
from neo4j.exceptions import CypherSyntaxError

from core.exception import CustomException

from crewplus.helper.graphutil import *

import logging

class Neo4jGraphService(object):
    """
    graph database service class

    """    
    def __init__(self, url, user, password, llm):
        self.driver = GraphDatabase.driver(url, auth=(user, password))

        self.llm = llm
        # construct schema
        self.schema = self.generate_schema()

    def generate_schema(self):
        logging.info("generate schema")
        node_props = self.query_database(node_properties_query)
        rel_props = self.query_database(rel_properties_query)
        rels = self.query_database(rel_query)
        return schema_text(node_props, rel_props, rels)

    def refresh_schema(self):
        self.schema = self.generate_schema()

    def get_system_message(self):
        return f"""
        Task: Generate Cypher queries to query a Neo4j graph database based on the provided schema definition.
        **Considerations for the Cypher Query:**
        * The query should be flexible enough to handle variations in phrasing, synonyms, and potentially incomplete information in the natural language question. * Use fuzzy matching techniques (e.g., `CONTAINS`) and case-insensitive comparisons where appropriate.

        Instructions:
        Use only the provided relationship types and properties.
        Do not use any other relationship types or properties that are not provided.
        Do not generate leading ```cypher or trailing ``` for cypher query
        If you cannot generate a Cypher statement based on the provided schema, try split user question into two parts in json format,
        first key is 'cypher', value is Cypher query with provided schema and second part key is 'question', value is the rest of user question.

        Please try your best to get root causes' relevant documents also, even if the user question does not require document or reference.
        
        **Example:**
        MATCH (r:Reason)-[:caused_by]->(rc:RootCause)
        WHERE toLower(r.reasoncode_description) CONTAINS toLower("Mechanical Failures")
        OPTIONAL MATCH (rc)-[:has_document]->(doc:CP_Document)
        RETURN rc.rootcause_description AS RootCause, doc.filename AS DocumentFilename, doc.path AS DocumentPath, doc.category AS DocumentCategory

        **Example:**
        MATCH rc:RootCause
        WHERE toLower(rc.rootcause_description) CONTAINS toLower("Lubrication")
        OPTIONAL MATCH (rc)-[:has_document]->(doc:CP_Document)
        RETURN rc.rootcause_description AS RootCause, doc.filename AS DocumentFilename, doc.path AS DocumentPath, doc.category AS DocumentCategory

        Schema:
        {self.schema}

        Note: Do not include any explanations or apologies in your responses.
        """

    def query_database(self, neo4j_query, params={}, question = ''):
        logging.info("query_database")
        try:
            with self.driver.session() as session:
                result = session.run(neo4j_query, params)
                output = [r.values() for r in result]
                output.insert(0, result.keys())
                return output
        except CypherSyntaxError as e:
            logging.info("there is no neo4j result, use llm to generate neo4j query")
            # retry with syntax error message
            history = [
                {"role": "assistant", "content": neo4j_query},
                {
                    "role": "user",
                    "content": f"""This query returns an error: {str(e)}
                    Give me an improved query that works without any explanations or apologies""",
                },
            ]
            
            cypher = self.construct_cypher(question, history)
            
            try:
                with self.driver.session() as session:
                    result = session.run(cypher, params)
                    output = [r.values() for r in result]
                    output.insert(0, result.keys())
                    return output
            except CypherSyntaxError as e:
                raise CustomException(str(e), code=status.HTTP_ERROR)


    def construct_cypher(self, question, history=None):
        logging.info("construct——cypher")
        messages = [
            {"role": "system", "content": self.get_system_message()},
            {"role": "user", "content": question},
        ]
        # Used for Cypher healing flows
        if history:
            messages.extend(history)

        result = self.llm.invoke(messages)

        return result.content

    def run(self, question, history=None, retry=True):

        # Construct Cypher statement
        cypher = self.construct_cypher(question, history)
        logging.info(cypher)
        try:
            return self.query_database(cypher)
        # Self-healing flow
        except CypherSyntaxError as e:
            # If out of retries
            if not retry:
              return "Invalid Cypher syntax"
        # Self-healing Cypher flow by
        # providing specific error to GPT-4
            logging.info("Retrying..")
            return self.run(
                question,
                [
                    {"role": "assistant", "content": cypher},
                    {
                        "role": "user",
                        "content": f"""This query returns an error: {str(e)}
                        Give me an improved query that works without any explanations or apologies""",
                    },
                ],
                retry=False
            )